/* -------------------------------------------------------------------------- *
   TestWinUnit - Maria Blees (maria.blees@microsoft.com)

   File: ImportHookerTest

   This file contains tests for ImportHooker.  It uses SystemTimeToFileTime
   as the test function to be hooked and unhooked.
 * -------------------------------------------------------------------------- */
#include "WinUnit.h"
#include "DllMain.h"
#include "MockLogger.h"
#include "ImportHooker.h"
#include "NulFile.h"

using namespace WinUnitLib;

#pragma region Helper data for HookAndUnhook test

namespace
{
    const SYSTEMTIME s_fakeSystemTime = { 
            1975, 
            10, 
            1,      // Monday
            13,
            0,
            0,
            0,
            0
        };

    // This value was found using SystemTimeToFileTime on the above SYSTEMTIME.
    const FILETIME s_fakeFileTime = {
        0x9344c000,
        0x01a42cb3
    };

    // Always returns 10/13/1975
    BOOL WINAPI FakeSystemTimeToFileTime(
      const SYSTEMTIME* lpSystemTime,
      LPFILETIME lpFileTime
    )
    {
        if (lpFileTime == NULL || lpSystemTime == NULL) { return FALSE; }
        *lpFileTime = s_fakeFileTime;
        return TRUE;
    }

    bool operator==(const SYSTEMTIME& lhs, const SYSTEMTIME& rhs)
    {
        return (
            lhs.wDay == rhs.wDay &&
            lhs.wDayOfWeek == rhs.wDayOfWeek &&
            lhs.wHour == rhs.wHour &&
            lhs.wMilliseconds == rhs.wMilliseconds &&
            lhs.wMinute == rhs.wMinute &&
            lhs.wMonth == rhs.wMonth &&
            lhs.wSecond == rhs.wSecond &&
            lhs.wYear == rhs.wYear);
    }

    bool operator==(const FILETIME& lhs, const FILETIME& rhs)
    {
        return (
            lhs.dwHighDateTime == rhs.dwHighDateTime &&
            lhs.dwLowDateTime == rhs.dwLowDateTime);
    }

    BOOL WINAPI CallSystemTimeToFileTime(
      const SYSTEMTIME* lpSystemTime,
      LPFILETIME lpFileTime
    );

}

template<>
inline const TCHAR* WinUnit::ToString(const FILETIME& object, TCHAR buffer[], size_t bufferSize)
{
    SYSTEMTIME systemTime;
    ::FileTimeToSystemTime(&object, &systemTime);
    ::_sntprintf_s(buffer, bufferSize, _TRUNCATE, 
        _T("{%d/%d/%d %d:%d}"),
        systemTime.wMonth,
        systemTime.wDay,
        systemTime.wYear,
        systemTime.wHour,
        systemTime.wMinute);
    return buffer;
}

#pragma endregion

// This test hooks and unhooks SystemFileTime and verifies that the results
// are as expected.
BEGIN_TEST(ImportHookerHookAndUnhook)
{
    // The function we're going to be hooking is SystemTimeToFileTime.  Our
    // fake SystemTimeToFileTime always returns the file time for October 13,
    // 1975.
    const unsigned int numberOfFunctionsToHook = 1;

    // This is just some date that is NOT October 13, 1975.  Pre-hook, 
    // SystemTimeToFileTime will return the right FILETIME for this one. 
    // Post-hook, it will be 10/13/1975.
    const SYSTEMTIME testSystemTime = {
        2007, 
        3, 
        4,      // Thursday
        29,
        0,
        0,
        0,
        0
    };
    MockLogger logger;

    {
        // First verify that the function is not hooked
        //
        // ... Convert testTime to fileTime
        FILETIME testFileTime;
        WIN_ASSERT_TRUE(::SystemTimeToFileTime(&testSystemTime, &testFileTime));
        //
        // ... Change fileTime back to SYSTEMTIME and verify that it's the same as
        // testTime
        SYSTEMTIME roundtripTestSystemTime;
        WIN_ASSERT_TRUE(::FileTimeToSystemTime(&testFileTime, &roundtripTestSystemTime));
        WIN_ASSERT_EQUAL(testSystemTime, roundtripTestSystemTime);
    }

    PROC originalFunctions[numberOfFunctionsToHook] = { 0 };
    {
        // Hook the function
        DWORD numberOfFunctionsHooked = 0;
        ImportHooker::HOOKFUNCDESC hookArray[numberOfFunctionsToHook] =
        {
            { "SystemTimeToFileTime", (PROC)FakeSystemTimeToFileTime }
        };

        WIN_ASSERT_TRUE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            "kernel32.dll", 
            numberOfFunctionsToHook, 
            hookArray, 
            originalFunctions, 
            &numberOfFunctionsHooked, 
            logger));
        WIN_ASSERT_EQUAL(numberOfFunctionsToHook, numberOfFunctionsHooked);
    }

    {
        // Verify that it is hooked
        FILETIME postHookFileTime;
//        WIN_ASSERT_TRUE(::SystemTimeToFileTime(&testSystemTime, &postHookFileTime));
        WIN_ASSERT_TRUE(CallSystemTimeToFileTime(&testSystemTime, &postHookFileTime));

        // This line fails in Release builds!
        WIN_ASSERT_EQUAL(s_fakeFileTime, postHookFileTime);
    }

    {
        // Change it back again
        ImportHooker::HOOKFUNCDESC hookArray[numberOfFunctionsToHook] =
        {
            { "SystemTimeToFileTime", originalFunctions[0] }
        };
        WIN_ASSERT_TRUE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            "kernel32.dll",
            numberOfFunctionsToHook,
            hookArray,
            NULL,
            NULL,
            logger));
    }

    {
        // Verify it's no longer hooked
        FILETIME postUnhookFileTime;
        WIN_ASSERT_TRUE(::SystemTimeToFileTime(&testSystemTime, &postUnhookFileTime));

        SYSTEMTIME roundtripTestSystemTime;
        WIN_ASSERT_TRUE(::FileTimeToSystemTime(&postUnhookFileTime, &roundtripTestSystemTime));
        WIN_ASSERT_EQUAL(testSystemTime, roundtripTestSystemTime);
    }
}
END_TEST

BEGIN_TEST(ImportHookerBadOrigFuncs)
{
    MockLogger logger;

    const ImportHooker::HOOKFUNCDESC hookArray[1] = {
        { "SystemTimeToFileTime", NULL }
    };

    WIN_ASSERT_FALSE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            "kernel32.dll", 
            1, 
            hookArray, 
            (PROC*)0x00000001, 
            NULL, 
            logger));

    WIN_ASSERT_EQUAL(1, logger.OutputProgramErrorLineCalled);
} 
END_TEST

BEGIN_TEST(ImportHookerBadPdwHookedParam)
{
    MockLogger logger;
    ImportHooker::HOOKFUNCDESC hookArray[1];

    WIN_ASSERT_FALSE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            "kernel32.dll", 
            1, 
            hookArray, 
            NULL, 
            (DWORD*)0x00000001,
            logger));

    WIN_ASSERT_EQUAL(1, logger.OutputProgramErrorLineCalled);
} 
END_TEST

BEGIN_TEST(ImportHookerBadHModule)
{
    MockLogger logger;

    ImportHooker::HOOKFUNCDESC hookArray[1];
    WIN_ASSERT_FALSE(ImportHooker::HookImportedFunctionsByName(
            0,
            "kernel32.dll", 
            1, 
            hookArray, 
            NULL, 
            NULL, 
            logger));
    WIN_ASSERT_EQUAL(1, logger.OutputProgramErrorLineCalled);
} 
END_TEST

BEGIN_TEST(ImportHookerBadImportModuleName)
{
    MockLogger logger;

    ImportHooker::HOOKFUNCDESC hookArray[1];
    WIN_ASSERT_FALSE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            NULL, 
            1, 
            hookArray, 
            NULL, 
            NULL, 
            logger));
    WIN_ASSERT_EQUAL(1, logger.OutputProgramErrorLineCalled);
} 
END_TEST

BEGIN_TEST(ImportHookerNullHookArray)
{
    MockLogger logger;

    WIN_ASSERT_FALSE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            "kernel32.dll", 
            1, 
            NULL, 
            NULL, 
            NULL, 
            logger));
    WIN_ASSERT_EQUAL(1, logger.OutputProgramErrorLineCalled);
} 
END_TEST

BEGIN_TEST(ImportHookerBadHookArray)
{
    MockLogger logger;

    WIN_ASSERT_FALSE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            "kernel32.dll", 
            1, 
            (ImportHooker::HOOKFUNCDESC*)1, 
            NULL, 
            NULL, 
            logger));
    WIN_ASSERT_EQUAL(1, logger.OutputProgramErrorLineCalled);
} 
END_TEST

BEGIN_TEST(ImportHookerBadFunctionNameInHookArray)
{
    MockLogger logger;

    ImportHooker::HOOKFUNCDESC hookArray[1] = {
        { NULL, (PROC)FakeSystemTimeToFileTime }
    };
    WIN_ASSERT_FALSE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            "kernel32.dll", 
            1, 
            hookArray, 
            NULL, 
            NULL, 
            logger));
    WIN_ASSERT_EQUAL(1, logger.OutputProgramErrorLineCalled);
} 
END_TEST

BEGIN_TEST(ImportHookerBadProcInHookArray)
{
    MockLogger logger;

    ImportHooker::HOOKFUNCDESC hookArray[1] = {
        { "SystemTimeToFileTime", (PROC)1 }
    };
    WIN_ASSERT_FALSE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            "kernel32.dll", 
            1, 
            hookArray, 
            NULL, 
            NULL, 
            logger));
    WIN_ASSERT_EQUAL(1, logger.OutputProgramErrorLineCalled);
} 
END_TEST

BEGIN_TEST(ImportHookerUiCountZero)
{
    MockLogger logger;

    ImportHooker::HOOKFUNCDESC hookArray[1] = {
        { "SystemTimeToFileTime", (PROC)FakeSystemTimeToFileTime }
    };
    WIN_ASSERT_TRUE(ImportHooker::HookImportedFunctionsByName(
            g_hmodule,
            "kernel32.dll", 
            0, 
            hookArray, 
            NULL, 
            NULL, 
            logger));
    WIN_ASSERT_EQUAL(0, logger.OutputProgramErrorLineCalled);
} 
END_TEST

// Change by JER. In an optimized release build, the compiler gets very 
// aggressive and will inline this function away. Additionally, it will optimize 
// the call to SystemTimeToFileTime to the direct export call so the hooking 
// is not called. By disabling optimizations on only this function, the above 
// tests can be build fully optimized to ensure they work.
#pragma optimize ( "" , off )
namespace
{
    BOOL WINAPI CallSystemTimeToFileTime(
      const SYSTEMTIME* lpSystemTime,
      LPFILETIME lpFileTime
    )
    {
        return ::SystemTimeToFileTime(lpSystemTime, lpFileTime);
    }
}
#pragma optimize ( "" , on )
